import os
from jupyter_client import KernelManager, find_connection_file
from queue import Empty
import re


def clean_traceback(traceback_list):
    cleaned_traceback = []
    for line in traceback_list:
        # 去掉ANSI转义字符
        clean_line = re.sub(r'\x1b\[.*?m', '', line)
        cleaned_traceback.append(clean_line)
    res = '\n'.join(cleaned_traceback)
    if len(res) > 1000:
        res = res[:500] + '...' + res[-500:]
    return res


class CodeKernel:
    def __init__(self):
        self.km = None
        self.kc = None

    def start(self):
        """Start the IPython kernel."""
        if not self.km or not self.km.is_alive():
            # Start a new kernel manager and connect to it
            self.km = KernelManager(kernel_name='my_kernel')
            self.km.start_kernel()
            self.kc = self.km.client()
            self.kc.start_channels()

            # Wait for the kernel to be ready
            self.kc.wait_for_ready()

    def shutdown(self):
        """Shutdown the IPython kernel."""
        if self.km and self.km.is_alive():
            self.kc.stop_channels()
            self.km.shutdown_kernel(now=True)
            self.km = None
            self.kc = None

    def run_stream(self, code):
        if not self.km or not self.km.is_alive():
            raise RuntimeError("Kernel is not running. Call start() first.")

        # Execute the code and get the message ID
        msg_id = self.kc.execute(code)
        print_info = ''
        # while '这是我自己添加的结束标志【23333333】' not in print_info:
        while 1:
            try:
                msg = self.kc.get_iopub_msg(timeout=2)
                msg_type = msg['msg_type']
                content = msg['content']

                if msg_type == 'stream':
                    yield f"{content['text']}\n\n"
                elif msg_type == 'execute_result' or msg_type == 'display_data':
                    if 'data' in content and 'text/plain' in content['data']:
                        print_info = content['data']['text/plain']
                        yield f"{content['data']['text/plain']}\n\n"

                elif msg_type == 'error':
                    yield f"Error: {''.join(content['traceback'])}\n\n"
                    break
                elif msg_type == 'status' and content.get('execution_state') == 'idle':
                    # Kernel is idle, but we need to confirm execution is done
                    pass
                elif msg_type == 'execute_reply' and msg['parent_header'].get('msg_id') == msg_id:
                    # Check if this is the reply for our executed code
                    if content['status'] == 'ok':
                        break  # Execution completed successfully
                    elif content['status'] == 'error':
                        yield f"Error: {''.join(content['traceback'])}\n\n"
                        break  # Execution completed with an error

            except Empty:
                # No more messages, assume execution is done
                break

    def run(self, code):
        """Execute Python code in the IPython kernel and return all output."""
        if not self.km or not self.km.is_alive():
            raise RuntimeError("Kernel is not running. Call start() first.")

        # Execute the code
        self.kc.execute(code)

        # Collect all outputs
        outputs = []
        while True:
            try:
                msg = self.kc.get_iopub_msg(timeout=200)
                msg_type = msg['msg_type']
                content = msg['content']

                if msg_type == 'stream':
                    outputs.append(content['text'])
                elif msg_type == 'execute_result' or msg_type == 'display_data':
                    if 'data' in content and 'text/plain' in content['data']:
                        outputs.append(content['data']['text/plain'])
                elif msg_type == 'error':
                    outputs.append('Error: ')
                    outputs.append(clean_traceback(content['traceback']))
                elif msg_type == 'status' and content.get('execution_state') == 'idle':
                    # Kernel is idle, we've received all output
                    break
            except Empty:
                # No more messages, assume execution is done
                break

        return '\n'.join(outputs).strip()

    def __del__(self):
        try:
            self.shutdown()
        except:
            ...


# 使用示例
if __name__ == "__main__":
    for i in range(100):
        kernel = CodeKernel()
        try:
            kernel.start()
            result = kernel.run("print('Hello, world!')\n1/0")
            print("Output from kernel:", result)
        finally:
            kernel.shutdown()
